# Treatment Effects by Imputation Rule ====
covs <- c("ln_capex", "D_ln_capex_missing", "boi_year", "D_boi_year_missing",
          "bh_ope_cost", "D_bh_ope_cost_missing", "ln_plant_total_heatoutput")

load_emissions_data <- function(EMISSIONS_DATA_OUT, covs) {
  
  RuleA_Panel <- read_dta(
    str_c(EMISSIONS_DATA_OUT, "/RuleA_Panel.dta"),
    col_select = all_of(c("lnY_rule0", "D_treatment", "gpcb_id", "month16", "D_interregnum", 
                          "post_mock1", "post_mock2", covs))) %>% 
    filter(D_interregnum == 0 & (post_mock1 == 1 | post_mock2 == 1)) %>% 
    select(-D_interregnum, -matches("post")) %>% 
    rename(lnY_rule_A_0 = lnY_rule0)
  
  RuleM_Panel <- read_dta(
    str_c(EMISSIONS_DATA_OUT, "/RuleM_Panel.dta"),
    col_select = c(matches("lnY_rule.*0$"),
      all_of(c("D_treatment", "gpcb_id", "month16", "D_interregnum", 
               "post_mock1", "post_mock2", covs)))) %>% 
    filter(D_interregnum == 0 & (post_mock1 == 1 | post_mock2 == 1)) %>% 
    select(-D_interregnum, -matches("post"))
  
  emissions_data <- RuleA_Panel %>% 
    full_join(RuleM_Panel, by = c("gpcb_id", "month16", "D_treatment", covs))
  
  # Confirm merge by confirming no missing values
  stopifnot(sum(is.na(emissions_data)) == 0)
  
  return(emissions_data)
}

get_TEs <- function(EMISSIONS_DATA_OUT, covs) {
  
  emissions_data <- load_emissions_data(EMISSIONS_DATA_OUT, covs)
  imputation_rules <- emissions_data %>% 
    select(matches("lnY")) %>% 
    names()
    
  TEs <- expand_grid(treated_rule = imputation_rules, control_rule = imputation_rules) %>% 
    mutate(
      data = map2(
        treated_rule, control_rule, 
        \(x, y) emissions_data %>% 
          mutate(
            # Set imputation according to appropriate rule
            yvar = if_else(D_treatment == 1, emissions_data[[x]], emissions_data[[y]]))),
      model = map(
        data, 
        \(x) fixest::feols(
          as.formula(str_c("yvar ~ D_treatment +", str_c(covs, collapse = " + "))), 
          data = x, vcov = ~ gpcb_id
        ) %>% 
          clean_model()
      )
    ) %>% 
    select(-data) %>% 
    unnest(everything())
  
  return(TEs)
  
}

make_table <- function(EMISSIONS_DATA_OUT, covs, EMISSIONS_TABS) {
  TEs <- get_TEs(EMISSIONS_DATA_OUT, covs)
  imputation_rules <- TEs %>% 
    pull(treated_rule) %>%
    unique()
  
  table_imputation_rules <- 
    c("lnY_rule_A_0", str_c("lnY_rule_p", c(70, 80, 90), "_0"), "lnY_rule_mkt_0")
  
  table_TEs <- TEs %>% 
    filter(treated_rule %in% table_imputation_rules & control_rule %in% table_imputation_rules)
  
  table <- map(table_imputation_rules, 
      \(x) filter(table_TEs, control_rule == x) %>% 
        mutate(order = map_int(
          treated_rule, ~which(.x == table_imputation_rules))) %>% 
        arrange(order, name) %>% 
        select(-order, - control_rule) %>% 
        rename_with(~x, .cols = value)) %>% 
    reduce(full_join, by = c("treated_rule", "name")) %>% 
    mutate(
      treated_rule = if_else(
      name == "estimate", 
      str_c("(", (row_number() + 1) / 2, ") ", map_chr(treated_rule, clean_names)),
      "")) %>% 
    select(-name) %>% 
    rename_with(clean_names) %>% 
    mutate(blank = "", .before = everything()) %>% 
    kbl(
      format = "latex",
      booktabs = T,
      escape = F,
      align = c("c","l","c","c","c","c","c"),
      linesep = "",
      col.names = NULL,
      toprule = 
        "\\toprule
          && \\multicolumn{5}{c}{Imputation rule -- control} \\\\
          \\cmidrule{3-7}
          && \\multicolumn{1}{c}{\\makecell{Rule A \\\\ (1)}} & \\multicolumn{1}{c}{\\makecell{p70 \\\\ (2)}} & \\multicolumn{1}{c}{\\makecell{p80 \\\\ (3)}} & \\multicolumn{1}{c}{\\makecell{p90 \\\\ (4)}} & \\multicolumn{1}{c}{\\makecell{Market \\\\ (5)}} \\\\
          \\cmidrule{2-7}
          \\multirow{4}{0.2in}{\\STAB{\\rotatebox[origin=c]{90}{ {\\fontsize{13}{14}\\selectfont Imputation rule -- treatment}\\hspace{0.2cm}}}} 
          \\multirow{4}{*}{\\STAB{\\rotatebox[origin=c]{90}{\\noindent\\rule{5cm}{0.1pt}\\hspace{0.2cm}}}} &&&&&\\\\[-0.3cm]",
      )
  
  save_kable(table, str_c(EMISSIONS_TABS, "/Table_C4.tex"))
  
  return(table)
}


## Helper functions ====
clean_model <- function(model) {
  
  fixest::coeftable(model)["D_treatment", c("Estimate", "Std. Error", "Pr(>|t|)")] %>% 
    tibble(name = c("estimate", "se", "pval"), value = .) %>% 
    mutate(
      pval = pull(., "value")[pull(., "name") == "pval"],
      value = 
        if_else(
          name == "estimate",
          str_c(
            formatC(value, digits = 3, format = "f"), 
            case_when(
              pval < 0.01 ~ "***",
              pval < 0.05 ~ "**",
              pval < 0.1 ~ "***",
              T ~ "")
          ),
          str_c("(", formatC(value, digits = 3, format = "f"), ")"))) %>% 
    filter(name != "pval") %>% 
    select(name, value) %>% 
    return()
}

clean_names <- function(name) {
  case_when(
    name == "lnY_rule_A_0" ~ "Rule A",
    name == "lnY_rule_mkt_0" ~ "Market",
    str_detect(name, "p\\d{2}") ~ str_extract(name, "p\\d{2}"), 
    T ~ " "
  ) %>% 
    return()
}

## Run ====
make_table(EMISSIONS_DATA_OUT, covs, EMISSIONS_TABS)